import json
import networkx as nx
from dao.DBdao import insertGraph


def constructDependencyGraph():
    trace_data = readTrace()
    nodes, edges = constructNodesAndEdges(trace_data)
    # graph = nx.DiGraph()
    # graph.add_nodes_from(nodes)
    # graph.add_edges_from(edges)
    insertGraph(nodes, edges)
    # return graph


def constructNodesAndEdges(trace_data):
    nodes = set()
    call = set()
    span_map = {}
    edges = set()
    for span in trace_data:
        nodes.add(span['serviceName'])
        span_map[span['spanId']] = span['serviceName']
        if span['parentId'] != -1 and span['parentId'] != span['spanId']:
            call.add((span['parentId'], span['spanId']))
    for item in call:
        if span_map[item[0]] != span_map[item[1]]:
            edges.add((span_map[item[0]], span_map[item[1]]))
    return list(nodes), list(edges)


def readTrace():
    f = open('../data/traces.json', 'r')
    original_traces = json.load(f)
    traces = {}
    for i, (k, v) in enumerate(original_traces.items()):
        if i <= 500:
            traces[k] = v
        else:
            break
    f.close()
    trace_data = list()
    for traceId, trace in traces.items():
        for index, span in enumerate(trace):
            if span['pid'] == 'None':
                span['pid'] = -1
            tmp_dict = {'spanId': span['id'], 'parentId': span['pid'], 'serviceId': span['serviceId'],
                        'serviceName': span['serviceName'], 'traceId': traceId}
            trace_data.append(tmp_dict)
    return trace_data


if __name__ == '__main__':
    constructDependencyGraph()
